package algorithm.graph.adjacencymatrix;

import algorithm.graph.base.Edge;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

public class Kruskal {
    AdjacencyMatrixGraph graph;

    public Kruskal(AdjacencyMatrixGraph graph) {
        this.graph = graph;
    }

    AdjacencyMatrixGraph buildMinEdgeGraph() {
        List<Edge> edgeList = buildMinEdgeGraphEdgeList(graph);
        AdjacencyMatrixGraph result = new AdjacencyMatrixGraph(graph.getNodeCount());
        result.initNonDirectedGraphEdges(edgeList);
        return result;
    }

    private static List<Edge> buildMinEdgeGraphEdgeList(AdjacencyMatrixGraph graph) {
        List<Integer> visitedNodeIdList = new ArrayList<>();
        List<Edge> unvisitedEdgeList = new ArrayList<>();
        List<Edge> validEdgeList = new ArrayList<>();
        int nodeCount = graph.getNodeCount();
        int[][] graphData = graph.getGraphData();
        for (int i = 0; i < nodeCount; i++) {
            for (int j = i + 1; j < nodeCount; j++) {
                int weight = graphData[i][j];
                if (weight != -1) {
                    Edge edge = new Edge(i, j, weight);
                    unvisitedEdgeList.add(edge);
                }
            }
        }
        while (!unvisitedEdgeList.isEmpty()) {
            Collections.sort(unvisitedEdgeList, (o1, o2) -> o1.weight() - o2.weight());
            Edge toRemoveEdge = unvisitedEdgeList.remove(0);
            if (visitedNodeIdList.contains(toRemoveEdge.from()) && visitedNodeIdList.contains(toRemoveEdge.to())) {
                continue;
            } else if (!visitedNodeIdList.contains(toRemoveEdge.from()) && !visitedNodeIdList.contains(toRemoveEdge.to())) {
                visitedNodeIdList.add(toRemoveEdge.from());
                visitedNodeIdList.add(toRemoveEdge.to());
                validEdgeList.add(toRemoveEdge);
            } else if (!visitedNodeIdList.contains(toRemoveEdge.from())) {
                visitedNodeIdList.add(toRemoveEdge.from());
                validEdgeList.add(toRemoveEdge);
            } else if (!visitedNodeIdList.contains(toRemoveEdge.to())) {
                visitedNodeIdList.add(toRemoveEdge.to());
                validEdgeList.add(toRemoveEdge);
            }
        }
        if (visitedNodeIdList.size() != nodeCount) {
            throw new RuntimeException("invalid graph");
        }
        return validEdgeList;
    }
}
